import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

def Lorenz(y0, t, para=[10,28,8/3]):
    # 设置默认参数可以不使用odeint的args参数
    p, r, b = para
    x, y, z = y0
    dx = -p*(x-y)
    dy = r*x-y-x*z
    dz = -b*z+x*y
    return np.array([dx,dy,dz])

t = np.arange(0, 60, 0.01)
y0 = [1,0,-1]
sol = odeint(Lorenz, y0, t)
# sol = odeint(Lorenz,y0,t,args=([10,28,8/3],)) # 定义函数时不使用默认参数则需要使用这个方法传参

plt.figure()
plt.axes(projection='3d')
plt.plot(sol[:, 0], sol[:, 1], sol[:, 2])


plt.figure(2)
plt.plot(t,sol)
plt.legend(['x','y','z'])
plt.xlabel('time')
plt.title('Lorenz')
plt.show()
sol = sol.T

np.save("Lorenz.npy", sol)